from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from typing import List


# llm = OllamaLLM(
#     model="qwen3:14b",  # 与 ollama list 中的模型名一致
#     base_url="http://localhost:11434",  # Ollama 服务地址
#     temperature=0.3,          # 控制随机性 (0-1)
#     num_predict=2048,         # 最大生成长度
#     top_k=50,                 # 采样参数
#     top_p=0.9,                # 采样参数
#     stop=["<|EOT|>"],         # DeepSeek 的停止标记
#     num_gpu=99,               # 使用全部 GPU 层加速
#     num_thread=8,             # CPU 线程数 (4090 可设 8-12)
# )

from langchain_community.llms import Tongyi
from dotenv import load_dotenv
import os

# 加载环境变量
load_dotenv()

# 获取环境变量
api_key = os.getenv("DASHSCOPE_API_KEY")

# 初始化模型
llm = Tongyi(
    model="qwen-turbo",
    api_key=api_key,
    temperature=0.5,
    max_tokens=1024,
    max_retries=2
)

# 1. 配置JSON输出解析器
output_parser = JsonOutputParser()

# 2. 构建提示模板 (整合系统指令与用户输入)
system_template = SystemMessagePromptTemplate.from_template(
    "你是一名资深抖音内容策划专家，擅长用{style}风格创作爆款短视频。"
    "严格按以下模块生成脚本：\n"
    "1. 黄金开头（Hook）：5秒内吸引眼球的强冲击力内容\n"
    "2. 背景铺垫（Setup）：交代场景和问题，引发共鸣\n"
    "3. 核心内容（Core）：分步骤或分要点呈现核心信息（至少3点）\n"
    "4. 情绪升华（Emotion）：激发用户情感共鸣（如感动、好奇、愤怒）\n"
    "5. 行动引导（CTA）：明确引导用户点赞、关注或评论\n"
    "6. 结尾彩蛋（Easter Egg）：意外反转或隐藏福利，提升完播率\n\n"
    "重要：只返回JSON格式，不要包含任何思考过程、解释或其他文字！\n"
    "JSON格式：\n"
    "{{\n"
    '  \"hook\": \"黄金开头内容\",\n'
    '  \"setup\": \"背景铺垫内容\",\n'
    '  \"core\": [\"核心要点1\", \"核心要点2\", \"核心要点3\"],\n'
    '  \"emotion\": \"情绪升华内容\",\n'
    '  \"cta\": \"行动引导内容\",\n'
    '  \"easter_egg\": \"结尾彩蛋内容\"\n'
    "}}"
)
prompt = ChatPromptTemplate.from_messages([
    system_template,
    ("human", "主题方向：{topic}")
])

# 3. 组装执行链
chain = prompt | llm | output_parser

# 4. 调用示例
result = chain.invoke({
    "topic": "职场塑料姐妹花",
    "style": "犀利幽默"
})

print(result)

# print("JSON格式输出结果：")
# print(f"类型: {type(result)}")
# print(f"内容: {result}")

# # 5. 访问JSON中的具体字段
# print("\n=== 脚本内容详情 ===")
# print(f"黄金开头: {result.get('hook', 'N/A')}")
# print(f"背景铺垫: {result.get('setup', 'N/A')}")
# print(f"核心内容: {result.get('core', [])}")
# print(f"情绪升华: {result.get('emotion', 'N/A')}")
# print(f"行动引导: {result.get('cta', 'N/A')}")
# print(f"结尾彩蛋: {result.get('easter_egg', 'N/A')}")